import numpy as np
import torch



class SumTree(object):
    data_pointer = 0

    def __init__(self, capacity,device):
        self.capacity = capacity  # for all priority values
        self.tree = torch.zeros(2 * capacity - 1).to(device)
        self.data = torch.zeros(capacity, dtype=torch.int32).to(device)  # for all transitions
        self.device=device

    def add(self, p, data):
        tree_idx = self.data_pointer + self.capacity - 1
        self.data[self.data_pointer] = data  # store transition in self.data
        self.update(tree_idx, p)  # add p to the tree
        self.data_pointer += 1
        if self.data_pointer >= self.capacity:
            self.data_pointer = 0

    def update(self, tree_idx, p):
        change = p - self.tree[tree_idx]
        self.tree[tree_idx] = p
        while tree_idx != 0:
            tree_idx = (tree_idx - 1) // 2
            self.tree[tree_idx] += change

    def get_leaf(self, v):
        parent_idx = 0
        while True:
            cl_idx = 2 * parent_idx + 1  # left kid of the parent node
            cr_idx = cl_idx + 1
            if cl_idx >= len(self.tree):  # kid node is out of the tree, so parent is the leaf node
                leaf_idx = parent_idx
                break
            else:  # downward search, always search for a higher priority node
                if v <= self.tree[cl_idx]:
                    parent_idx = cl_idx
                else:
                    v -= self.tree[cl_idx]
                    parent_idx = cr_idx

        data_idx = leaf_idx - self.capacity + 1
        return leaf_idx, self.tree[leaf_idx], self.data[data_idx]

    @property
    def total_p(self):
        return self.tree[0].to('cpu')